{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BfNNHaU1ouzC"
   },
   "source": [
    "# ML Simulation Evaluation\n",
    "In this notebook you are going to control both the SDV and other agents using a CNN-based policy in a closed-loop fashion.\n",
    "\n",
    "**Note: to learn more about closed-loop evaluation refer to our [planning notebook](../planning/closed_loop_test.ipynb).**\n",
    "\n",
    "![simulation-example](https://github.com/lyft/l5kit/blob/master/docs/images/simulation/simulation_example.svg?raw=1)\n",
    "\n",
    "\n",
    "## What can we use ML simulation for?\n",
    "Simulating other agents **is crucial to remove false positive interventions**, which occur when agents are log-replayed.\n",
    "As an example, imagine if our SDV was slower compared to the one in the data log. In this situation the car behind us may bump into us if it's just being replayed from the log. Differently, if that agent is equipped with a ML policy it will be able to react and slow down before colliding with the SDV.\n",
    "\n",
    "When evaluating a policy for the SDV, ML simulation can therefore help in reducing the noise introduced in the metrics by the non-reactive log replayed agents. \n",
    "\n",
    "![simulation-bump](https://github.com/lyft/l5kit/blob/master/docs/images/simulation/simulation_bump.svg?raw=1)\n",
    "\n",
    "## Why is ML simulation difficult?\n",
    "Clearly, the policy from the previous scenario must be smart enough to model a chasing car that is:\n",
    "- **realistic**: such that a simulated agent should be indistinguishable from a log-replayed one in standard conditions;\n",
    "- **reactive**: such that a simulated agent can react to the SDV and other agents around it in the same way a real agents would do on the road. \n",
    "\n",
    "In practice, **solving ML simulation is at least as hard as solving planning for the SDV itself**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-0BBPxheouzE"
   },
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "LKI_5lbiouzF"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from l5kit.configs import load_config_data\n",
    "from l5kit.data import LocalDataManager, ChunkedDataset, filter_agents_by_frames\n",
    "from l5kit.dataset import EgoDataset\n",
    "from l5kit.rasterization import build_rasterizer\n",
    "\n",
    "from l5kit.cle.closed_loop_evaluator import ClosedLoopEvaluator, EvaluationPlan\n",
    "from l5kit.cle.metrics import (CollisionFrontMetric, CollisionRearMetric, CollisionSideMetric,\n",
    "                               DisplacementErrorL2Metric, DistanceToRefTrajectoryMetric)\n",
    "from l5kit.cle.validators import RangeValidator, ValidationCountingAggregator\n",
    "from l5kit.simulation.dataset import SimulationConfig\n",
    "from l5kit.simulation.unroll import ClosedLoopSimulator\n",
    "import os\n",
    "from collections import defaultdict\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from prettytable import PrettyTable\n",
    "from l5kit.visualization.visualizer.zarr_utils import simulation_out_to_visualizer_scene\n",
    "from l5kit.visualization.visualizer.visualizer import visualize\n",
    "from bokeh.io import output_notebook, show\n",
    "from l5kit.data import MapAPI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EtLmUJ1CouzF"
   },
   "source": [
    "## Prepare data path and load cfg\n",
    "\n",
    "By setting the `L5KIT_DATA_FOLDER` variable, we can point the script to the folder where the data lies.\n",
    "\n",
    "Then, we load our config file with relative paths and other configurations (rasteriser, training params...)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "_XeBMM3oouzG"
   },
   "outputs": [],
   "source": [
    "#@title Download L5 Sample Dataset and install L5Kit\n",
    "import os\n",
    "RunningInCOLAB = 'google.colab' in str(get_ipython())\n",
    "if RunningInCOLAB:\n",
    "    !wget https://raw.githubusercontent.com/lyft/l5kit/master/examples/setup_notebook_colab.sh -q\n",
    "    !sh ./setup_notebook_colab.sh\n",
    "    os.environ[\"L5KIT_DATA_FOLDER\"] = open(\"./dataset_dir.txt\", \"r\").read().strip()\n",
    "else:\n",
    "    print(\"Not running in Google Colab.\")\n",
    "    os.environ[\"L5KIT_DATA_FOLDER\"] = \"/tmp/l5kit_data\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EZI29omqouzG"
   },
   "outputs": [],
   "source": [
    "# set env variable for data\n",
    "dm = LocalDataManager(None)\n",
    "# get config\n",
    "cfg = load_config_data(\"./config.yaml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4i2TgS2douzH"
   },
   "source": [
    "## Load the model\n",
    "\n",
    "We load here two models:\n",
    "- the `simulation_model.pt` will be used to control the agents around the SDV;\n",
    "- the `ego_model.pt` will be used to control the SDV itself;\n",
    "\n",
    "Clearly, nothing prevents us from replacing one model with the other (or with a completely different one, provided the input and output stay the same)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vcEdI_WWouzH"
   },
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "simulation_model_path = \"/tmp/simulation_model.pt\"\n",
    "simulation_model = torch.load(simulation_model_path).to(device)\n",
    "simulation_model = simulation_model.eval()\n",
    "\n",
    "ego_model_path = \"/tmp/planning_model.pt\"\n",
    "ego_model = torch.load(ego_model_path).to(device)\n",
    "ego_model = ego_model.eval()\n",
    "\n",
    "torch.set_grad_enabled(False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "T9mfiTF4ouzH"
   },
   "source": [
    "## Load the evaluation data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kCwROryUouzI"
   },
   "outputs": [],
   "source": [
    "# ===== INIT DATASET\n",
    "eval_cfg = cfg[\"val_data_loader\"]\n",
    "rasterizer = build_rasterizer(cfg, dm)\n",
    "mapAPI = MapAPI.from_cfg(dm ,cfg)\n",
    "eval_zarr = ChunkedDataset(dm.require(eval_cfg[\"key\"])).open()\n",
    "eval_dataset = EgoDataset(cfg, eval_zarr, rasterizer)\n",
    "print(eval_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ImXyrGpUouzI"
   },
   "source": [
    "## Define some simulation properties\n",
    "We define here some common simulation properties such as the length of the simulation and how many scene to simulate.\n",
    "\n",
    "**NOTE: these properties have a significant impact on the execution time. We suggest you to increase them only if your setup includes a GPU**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "v1w1j-CdouzJ"
   },
   "outputs": [],
   "source": [
    "scenes_to_unroll = [0, 10, 20]\n",
    "num_simulation_step_example1 = 20\n",
    "num_simulation_step_example2 = 20"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-aykbmcpouzJ"
   },
   "source": [
    "# Measuring agents realism\n",
    "The first set of metric we're interested in answers to the question _\"how realistic are the simulated agents?\"_\n",
    "\n",
    "To this end, we can run a simulation with the following features:\n",
    "- agents only are simulated (we're not interested in the SDV for now);\n",
    "- only agents in the initial frames are simulated (this avoid interactions between log-replayed and simulated agents);\n",
    "- once an agent is simulated, it's simulated until the end of the episode (note `distance_th_far=500`).\n",
    "\n",
    "**Quantitatively**: we can measure the L2 displacement between the log-replayed and the simulated agents at each time step (we can report this in terms of FDE, ADE, etc..)\n",
    "\n",
    "**Qualitatively**: we can inspect the resulting simulation overlayed with the original trajectories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kluQWjkYouzJ"
   },
   "outputs": [],
   "source": [
    "sim_cfg = SimulationConfig(use_ego_gt=True, use_agents_gt=False, disable_new_agents=True,\n",
    "                           distance_th_far=500, distance_th_close=50, num_simulation_steps=num_simulation_step_example1,\n",
    "                           start_frame_index=0, show_info=True)\n",
    "\n",
    "sim_loop = ClosedLoopSimulator(sim_cfg, eval_dataset, device, model_agents=simulation_model)\n",
    "\n",
    "sim_outs = sim_loop.unroll(scenes_to_unroll)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dGaSpVShouzK"
   },
   "source": [
    "# Quantitative evaluation\n",
    "We can measure L2 displacements between annotated and simulated positions and report them in terms of per step ADE.\n",
    "This error reflects how far the simulated agents are from the annotated ones.\n",
    "\n",
    "Because the annotated agents come from realistic log-replays, **this value can be interpreted as a way of measuring how realistic the simulation is**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "N4UBtDEwouzK"
   },
   "outputs": [],
   "source": [
    "errors_at_step = defaultdict(list)\n",
    "for sim_out in sim_outs: # for each scene\n",
    "    for idx_step, agents_in_out in enumerate(sim_out.agents_ins_outs):  # for each step\n",
    "        for agent_in_out in agents_in_out:  # for each agent\n",
    "            annot_pos = agent_in_out.inputs[\"target_positions\"][0]\n",
    "            pred_pos = agent_in_out.outputs[\"positions\"][0]\n",
    "            if agent_in_out.inputs[\"target_availabilities\"][0] > 0:\n",
    "                errors_at_step[idx_step + 1].append(np.linalg.norm(pred_pos - annot_pos))\n",
    "\n",
    "time_steps = np.asarray(list(errors_at_step.keys()))\n",
    "errors = np.asarray([np.mean(errors_at_step[k]) for k in errors_at_step])\n",
    "plt.plot(time_steps, errors, label=\"per step ADE\")\n",
    "plt.xticks(time_steps)\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "A0rVzQ-couzK"
   },
   "source": [
    "# Qualitative evaluation\n",
    "We can show a visualisation of the scene using our bokeh-based visualzer.\n",
    "\n",
    "For each scene:\n",
    "- the simulated scene is played;\n",
    "- both replayed and simulated trajectory can be toggled.\n",
    "\n",
    "The visualisation is interactive, **try to hover over an agent to show more information about it**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "C222hTKFouzK",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "output_notebook()\n",
    "for sim_out in sim_outs: # for each scene\n",
    "    vis_in = simulation_out_to_visualizer_scene(sim_out, mapAPI)\n",
    "    show(visualize(sim_out.scene_id, vis_in))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8pCdkKN3ouzL"
   },
   "source": [
    "# Measuring ego metrics with simulated agents\n",
    "The second set of metric we're interested in answer to the question _\"are simulated agents helpful when evaluating SDV?\"_\n",
    "\n",
    "To this end, we can run a simulation with the following features:\n",
    "- both agents and the SDV are simulated (so we can measure interactions between them);\n",
    "- only agents close to the SDV are simulated (this is the only set we really care about in terms of interaction with the SDV);\n",
    "- the set of simulated agents is variable and depends on the proximity of SDV (as the SDV moves in the scene, the agents proximity changes).\n",
    "\n",
    "**Quantitatively**, we can measure how these agents affect the closed-loop metrics for the SDV by comparing the results with those from the log-replayed agents.\n",
    "\n",
    "**Qualitatively**, we can inspect the resulting simulation overlayed with the original trajectories."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "x3LTCMDhouzL"
   },
   "source": [
    "# Quantitative evaluation\n",
    "\n",
    "We define here which metrics are of interest and how to evaluate each one of them. These metrics applies to the SDV only and are the same used in the [closed-loop planning notebook](../planning/closed_loop_test.ipynb).\n",
    "\n",
    "If you want to know more about how each component (metrics, validators, aggregators) of the evaluation stack works, please refer to the source code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9mB8wdaaouzL"
   },
   "outputs": [],
   "source": [
    "metrics = [DisplacementErrorL2Metric(),\n",
    "           DistanceToRefTrajectoryMetric(),\n",
    "           CollisionFrontMetric(),\n",
    "           CollisionRearMetric(),\n",
    "           CollisionSideMetric()]\n",
    "\n",
    "validators = [RangeValidator(\"displacement_error_l2_validator\", DisplacementErrorL2Metric, max_value=30),\n",
    "              RangeValidator(\"distance_ref_trajectory_validator\", DistanceToRefTrajectoryMetric, max_value=4),\n",
    "              RangeValidator(\"collision_front_validator\", CollisionFrontMetric, max_value=0),\n",
    "              RangeValidator(\"collision_rear_validator\", CollisionRearMetric, max_value=0),\n",
    "              RangeValidator(\"collision_side_validator\", CollisionSideMetric, max_value=0),\n",
    "              ]\n",
    "\n",
    "intervention_validators = [\"displacement_error_l2_validator\",\n",
    "                           \"distance_ref_trajectory_validator\",\n",
    "                           \"collision_front_validator\",\n",
    "                           \"collision_rear_validator\",\n",
    "                           \"collision_side_validator\"]\n",
    "\n",
    "cle_evaluator = ClosedLoopEvaluator(EvaluationPlan(metrics=metrics,\n",
    "                                    validators=validators,\n",
    "                                    composite_metrics=[],\n",
    "                                    intervention_validators=intervention_validators))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GZ6dzLiJouzM"
   },
   "source": [
    "### Evaluating with simulated agents\n",
    "In this evaluation **both the models for planning and simulation are enabled**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6ru6ePJcouzM"
   },
   "outputs": [],
   "source": [
    "sim_cfg = SimulationConfig(use_ego_gt=False, use_agents_gt=False, disable_new_agents=False,\n",
    "                           distance_th_far=30, distance_th_close=15, num_simulation_steps=num_simulation_step_example2,\n",
    "                           start_frame_index=0, show_info=True)\n",
    "\n",
    "sim_loop = ClosedLoopSimulator(sim_cfg, eval_dataset, device, model_ego=ego_model, model_agents=simulation_model)\n",
    "\n",
    "sim_outs = sim_loop.unroll(scenes_to_unroll)\n",
    "\n",
    "cle_evaluator.evaluate(sim_outs)\n",
    "validation_results = cle_evaluator.validation_results()\n",
    "agg = ValidationCountingAggregator().aggregate(validation_results)\n",
    "cle_evaluator.reset()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3ki5trAOouzM"
   },
   "source": [
    "### Evaluating with log-replayed agents\n",
    "In this evaluation **only the model for planning is enabled**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3Vyagd1WouzM"
   },
   "outputs": [],
   "source": [
    "sim_cfg_log = SimulationConfig(use_ego_gt=False, use_agents_gt=True, disable_new_agents=False,\n",
    "                           distance_th_far=30, distance_th_close=15, num_simulation_steps=num_simulation_step_example2,\n",
    "                           start_frame_index=0, show_info=True)\n",
    "\n",
    "sim_loop_log = ClosedLoopSimulator(sim_cfg_log, eval_dataset, device, model_ego=ego_model)\n",
    "sim_outs_log = sim_loop_log.unroll(scenes_to_unroll)\n",
    "\n",
    "cle_evaluator.evaluate(sim_outs_log)\n",
    "validation_results_log = cle_evaluator.validation_results()\n",
    "agg_log = ValidationCountingAggregator().aggregate(validation_results_log)\n",
    "cle_evaluator.reset()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XM1om61JouzN"
   },
   "source": [
    "### Comparing the two runs\n",
    "We compare here the result of replacing log-replayed agents with simulated one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QqMtf8yBouzN"
   },
   "outputs": [],
   "source": [
    "fields = [\"metric\", \"log_replayed agents\", \"simulated agents\"]\n",
    "table = PrettyTable(field_names=fields)\n",
    "for metric_name in agg_log:\n",
    "    table.add_row([metric_name, agg_log[metric_name].item(), agg[metric_name].item()])\n",
    "print(table)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yNWNhH2LouzN"
   },
   "source": [
    "# Qualitative evaluation\n",
    "Similarly to the previous example, also here we can show an interactive plot of the simulated scene. Because also the SDV is controlled now, its trajectory are shown in the plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2p0wRQs1ouzN",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "output_notebook()\n",
    "for sim_out in sim_outs: # for each scene\n",
    "    vis_in = simulation_out_to_visualizer_scene(sim_out, mapAPI)\n",
    "    show(visualize(sim_out.scene_id, vis_in))"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "Copy of simluation_test.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
